from asyncio import FastChildWatcher
from code import interact
import os
import time
import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
#from isaacgym import gymtorch
from math import sqrt
import math
from sympy import false
import torch
import cv2

from draw import *

from pcgworker.PCGWorker import *

from wfc_env_dynamic import *

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.results_plotter import load_results, ts2xy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common import results_plotter
from stable_baselines3.common.results_plotter import load_results, ts2xy, plot_results
from stable_baselines3.common.noise import NormalActionNoise
from stable_baselines3.common.callbacks import BaseCallback

LOGDIR = "./training_logs"

m_env = CustomEnv(seed_pth = "seed.json", headless_ = False)
observation = m_env.reset()

print("num_envs:",m_env.num_envs)

action = -1

steps_ = 0
reset_id = 0

PCGWorker_ = PCGWorker(9,9)

while True:

    steps_ += 1
    
    observation, reward, done, info = m_env.step(action)

    if steps_ >= 100 and steps_ % 10 == 0:

        if reset_id >= m_env.num_matrix_envs:
            reset_id = 0

        seed_ = PCGWorker_.generate()
        m_env.set_landscape(reset_id,seed_)
        seed_ = PCGWorker_.generate()
        m_env.set_landscape(reset_id+1,seed_)
        seed_ = PCGWorker_.generate()
        m_env.set_landscape(reset_id+2,seed_)
        seed_ = PCGWorker_.generate()
        m_env.set_landscape(reset_id+3,seed_)

        # m_env.reset_all_landscape()

        reset_id += 4

    # print(observation)
    # print(reward)
    # print(done)
    # print(info)

    interaction = m_env.render()

    if interaction == -2:
        m_env.close()

    if interaction == 119:  # w
        action = 0
    elif interaction == 115:    # s
        action = 1
    elif interaction == 97:     # a
        action = 2
    elif interaction == 100:    # d
        action = 3
    elif interaction == 107:    # ccw
        action = 4
    elif interaction == 108:    # cw
        action = 5
    # elif interaction == 111:    # break
    #     action = 6
    elif interaction == 114:
        m_env.reset()
    else:
        action = -1

